package Test.负载均衡;

import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicInteger;

public class WeightRoundRobinDemo {
    // 定义一个全局计数器，每次调用时递增
    private static AtomicInteger atomicInteger = new AtomicInteger(0);
    // 定义一个服务器及其对应权重值的映射
    private static Map<String, Integer> serverMap = new TreeMap<>();
    // 记录所有服务器的总权重
    private static int totalWeight = 0;
 
    public static void main(String[] args) {
        serverMap.put("Server A", 4);
        serverMap.put("Server B", 3);
        serverMap.put("Server C", 1);
        // 计算所有服务器的总权重
        for (Map.Entry<String, Integer> entry : serverMap.entrySet()) {
            totalWeight += entry.getValue();
        }
        // 模拟四个请求
        System.out.println(weightRoundRobin());
        System.out.println(weightRoundRobin());
        System.out.println(weightRoundRobin());
        System.out.println(weightRoundRobin());
        System.out.println(weightRoundRobin());
        System.out.println(weightRoundRobin());
        System.out.println(weightRoundRobin());
        System.out.println(weightRoundRobin());
    }
 
    public static String weightRoundRobin() {
        // 获取服务器数量
        int serverCount = serverMap.size();
        // 如果没有可用服务器，返回null
        if (serverCount == 0) {
            return null;
        }
        // 为避免多线程环境中并发操作导致的错误，在方法内部执行锁操作
        synchronized (serverMap) {
            // 确定当前请求应转发到哪个服务器
            int currentServerIndex = atomicInteger.incrementAndGet() % totalWeight;
            // 遍历服务器列表并根据服务器权重值选择相应地址
            for (Map.Entry<String, Integer> entry : serverMap.entrySet()) {
                String serverAddress = entry.getKey();
                Integer weight = entry.getValue();
                currentServerIndex -= weight;
                if (currentServerIndex < 0) {
                    return serverAddress;
                }
            }
        }
        return null;
    }
}